from funasr import AutoModel
import os
import soundfile

path_asr  = 'tools/asr/models/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'
path_vad  = 'tools/asr/models/speech_fsmn_vad_zh-cn-16k-common-pytorch'
path_punc = 'tools/asr/models/punc_ct-transformer_zh-cn-common-vocab272727-pytorch'
path_asr  = path_asr  if os.path.exists(path_asr)  else "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"
path_vad  = path_vad  if os.path.exists(path_vad)  else "iic/speech_fsmn_vad_zh-cn-16k-common-pytorch"
path_punc = path_punc if os.path.exists(path_punc) else "iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch"

model = AutoModel(
    model               = path_asr,
    model_revision      = "v2.0.4",
    vad_model           = path_vad,
    vad_model_revision  = "v2.0.4",
    punc_model          = path_punc,
    punc_model_revision = "v2.0.4",
)

# 语音转文本
# wav_path = "../../data/1_3月1日(1)_(Vocals).mp3"
# res = model.generate(input=wav_path)
# print(res)

def asr(input_audio_file):
    # 语音转文本
    res = model.generate(input=input_audio_file)
    return res[0]['text']

def batch_asr(input_folder, output_file):
    # 判断输出文件是否存在，不存在则创建，存在则清空
    if os.path.exists(output_file):
        os.remove(output_file)
    # 遍历输入文件夹中的所有文件
    for file_name in os.listdir(input_folder):
        if file_name.endswith('.wav'):  # 只处理 WAV 文件，可以根据需要修改
            # 构建输入文件的完整路径
            input_file_path = os.path.join(input_folder, file_name)

            # 执行语音转文本操作
            text = asr(input_file_path)

            # 保存文本到输出文件 每一句语音对应一个文本，｜隔开
            with open(output_file, "a", encoding="utf-8") as f:
                f.write(f"{input_file_path}|{text}\n")

# batch_asr("../../data/lanying/vad", "../../data/lanying/asr_out")

